# -*- coding: utf-8 -*-
# @Time    : 2023/5/31 21:45
# @Author  : Z.J. Zhang
# @Email   : zijingzhang@mail.ustc.edu.cn
# @File    : sw_to_MAGIC_commands.py
# @Software: PyCharm
"""
将一系列只包含单个实体（solid body）的.sldprt文件转化为MAGIC中的建模指令
"""
import datetime as dt
import enum
import inspect
import os.path
import typing

import matplotlib
import numpy

from _logging import logger

matplotlib.use('tkagg')

import matplotlib.pyplot as plt
from stl import mesh
from _common import facets_to_area, sldprt_to_stl

MAGIC_MAX_N_CHAR_PER_LINE = 69  # MAGIC输入文件每行允许的最大字符数目


def _auto_space(cmd_in_1_line: str, N_SPACE=1):
    """
    在命令前加上空格（以便VSCode对自动生成的命令进行折叠）
    :param cmd_in_1_line:
    :return:
    """
    return N_SPACE * " " + cmd_in_1_line


def _index_of_last_space_before_N(text, N):
    s = text[:N + 1]
    return len(s) - 1 - s[::-1].find(' ')


def _longest_words_not_longer_than_N(text, N=MAGIC_MAX_N_CHAR_PER_LINE):
    """
    总长不超过N的（完整单词组成的）最长字符串
    :param text:
    :param N:
    :return:
    """
    i = _index_of_last_space_before_N(text, N)
    # print(i)
    return text[:i], text[i:]


def auto_format(cmd_in_1_line_: str):
    """
    自动格式化：
        自动在行首加空格；
        对于过长的命令，自动在空格处换行。
    :return:
    """
    cmd_to_process = _auto_space(cmd_in_1_line_)  # 待处理的命令
    cmds = ""
    while len(cmd_to_process) > MAGIC_MAX_N_CHAR_PER_LINE:
        short_cmd, cmd_to_process = _longest_words_not_longer_than_N(cmd_to_process, MAGIC_MAX_N_CHAR_PER_LINE)
        cmd_to_process = _auto_space(cmd_to_process)
        cmds += short_cmd + "\n"
    return cmds + cmd_to_process


def pts_to_MAGIC_AREA_POLYGONAL_cmd(pts: numpy.ndarray, facet_prefix, facet_index):
    assert pts.shape[-1] == 2
    obj_name_in_MAGIC = "%s.%d" % (facet_prefix, facet_index)
    cmd = "AREA %s POLYGONAL " % (obj_name_in_MAGIC)
    for i in range(len(pts) + 1):
        cmd += "%.6e, %.6e " % (*pts[i % len(pts)],)
    cmd = cmd[:-1] + ";\n"
    return cmd


class Axis(enum.Enum):
    x = 0
    y = 1
    z = 2


import pypinyin, string


class NameConverter:
    @staticmethod
    def make_valid_MAGIC_Objectname_(name):
        valid_chars = "_.%s%s" % (string.ascii_letters, string.digits)
        cs = []
        for c in name:
            if c in valid_chars:
                cs.append(c)
            else:
                cs.append("_")
        name = ''.join(cs)
        return name

    # 判断字符串是否包含中文
    @staticmethod
    def str_contain_chinese(str):
        for ch in str:
            if u'\u4e00' <= ch <= u'\u9fff':
                return True
        return False

    @staticmethod
    def convert_to_valid_object_name_in_MAGIC(name):
        if NameConverter.str_contain_chinese(name):
            name = ''.join(pypinyin.lazy_pinyin(name))
        return NameConverter.make_valid_MAGIC_Objectname_(name)

def get_min_and_max(x):
    return numpy.min(x),numpy.max(x)
def area_command(filename: str, ax=plt.gca(),  # direction_of_stretching:Axis= Axis.x,
                 axis_of_sketch: typing.Tuple[Axis, Axis] = (Axis.z, Axis.y),
                 container_of_important_dimensions:typing.Dict[str, typing.List[float]] = None):
    # filename = r"D:\MagicFiles\HPM\电子枪部分\swCATHODE.STL"
    scaling_factor = 1e-3

    mymesh = mesh.Mesh.from_file(filename)
    # pts = _choose_2_axis(mymesh.points.reshape((-1, 3, 3))) * scaling_factor

    # _index = [True] *3
    # _index[direction_of_stretching.value] = False
    pts = mymesh.points.reshape((-1, 3, 3))[:, :, [axis.value for axis in axis_of_sketch]] * scaling_factor
    # plt.figure(facets_to_area(pts))
    polygon_pts = numpy.array(facets_to_area(pts).exterior.xy).T  # 一堆facet拼成的形状

    polygon_prefix = NameConverter.convert_to_valid_object_name_in_MAGIC(os.path.splitext(os.path.split(filename)[1])[
                                                                             0]
                                                                         # input("输入MAGIC中生成多边形的名称前缀，如’FaceFromSW‘")  # "FaceFromSW" # facet转化后的名称前缀
                                                                         )
    copyrightinfo_begin = "! Those codes are auto-generated by %s\n! =========area: %s auto-generated code start=========\n! ==============%s=============" % (
        __file__, polygon_prefix, dt.datetime.now().strftime("%Y/%m/%d %H:%M:%S"))
    copyrightinfo_end = "! =========area: %s auto-generated code end=========\n" % (polygon_prefix)
    _MAGIC_modeling_cmds = ""
    N_facets = mymesh.points.shape[0]

    _MAGIC_modeling_cmds += auto_format(pts_to_MAGIC_AREA_POLYGONAL_cmd(polygon_pts, polygon_prefix, 0))
    _MAGIC_modeling_cmds += auto_format("%s.N_facet=%d;\n" % (polygon_prefix, 1))
    x1_min,x1_max = get_min_and_max(polygon_pts[:, 0])
    x2_min,x2_max = get_min_and_max(polygon_pts[:, 1])

    container_of_important_dimensions__ = {
        "X1": [x1_min,x1_max],
        "X2": [x2_min,x2_max],
    }
    _MAGIC_modeling_cmds += auto_format("%s.x1mn=%.6e;\n" % (polygon_prefix,x1_min))
    _MAGIC_modeling_cmds += auto_format("%s.x1mx=%.6e;\n" % (polygon_prefix, x1_max))
    _MAGIC_modeling_cmds += auto_format("%s.x2mn=%.6e;\n" % (polygon_prefix, x2_min))
    _MAGIC_modeling_cmds += auto_format("%s.x2mx=%.6e;\n" % (polygon_prefix, x2_max))

    MAGIC_modeling_cmds = "%s\n%s\n%s" % (copyrightinfo_begin, _MAGIC_modeling_cmds, copyrightinfo_end)
    if ax:
        ax.plot(*polygon_pts.T, label=polygon_prefix)

    # print(MAGIC_modeling_cmds)


    if container_of_important_dimensions is not None:
        for key in container_of_important_dimensions__:
            container_of_important_dimensions[key] =  container_of_important_dimensions.get(key,[])+    container_of_important_dimensions__[key]
    return MAGIC_modeling_cmds


def sldprt_to_MAGIC_commands(sldprts: typing.List[str], ax=plt.gca(),
                             axis_of_sketch: typing.Tuple[Axis, Axis] = (Axis.z, Axis.y),important_dimensions :dict= None,re_export_to_stl=True):
    stl_files = []
    for sldprt in sldprts:
        stl_files.append(sldprt_to_stl(sldprt,re_export_to_stl))

    cmds = ''
    for stl_file in stl_files:
        cmds += area_command(stl_file, ax, axis_of_sketch,important_dimensions)
    with open('temp.txt', 'w') as f:
        f.write(cmds)
        logger.info(cmds)
    return cmds


def generate_grid_commands_to_keep_important_dimensions(important_pos_unit_in_m,
                                                        # important_x2s_unit_in_m:numpy.ndarray,
                                                        dimension_name="x2",
                                                        size_of_unimportant_cell: float = 1e-3,
                                                        tol_unit_in_m: float = 1e-9
                                                        ):
    """
    生成MAGIC中的GRID相关命令，用于在生成的网格中保持结构的关键尺寸不发生变化。
    Args:
        important_pos_unit_in_m:
        dimention_name:

    Returns:

    """

    important_pos_unit_in_m = numpy.sort(
        list(set((numpy.array(important_pos_unit_in_m) // tol_unit_in_m).astype(int)))) * tol_unit_in_m
    grid_origin = min(important_pos_unit_in_m)
    diff_important_pos_unit_in_m = numpy.diff(important_pos_unit_in_m)
    copyrightinfo_begin = "! Those codes are auto-generated by %s [function name: %s] \n! =========auto-generated code start=========\n! ==============%s=============" % (
        __file__,inspect.getframeinfo(inspect.currentframe()).function, dt.datetime.now().strftime("%Y/%m/%d %H:%M:%S"))
    mm = 1e-3

    comment = """
! important dimensions in %s = 
! [%s] mm    
    """%(dimension_name,_longest_words_not_longer_than_N(''.join(["%.8f, "%pos for pos in(  important_pos_unit_in_m /mm)]),MAGIC_MAX_N_CHAR_PER_LINE))
    # 注意，请不要修改下面字符串的缩进！！！！！！
    cmds = """
    
%s
%s

    GRID  ORIGIN %s %.8fmm;""" % (copyrightinfo_begin,comment, dimension_name, grid_origin / mm)
    cmds += "".join(["""
    GRID UNIFORM %s DISTANCE %.6fmm CELLS %d;""" % (dimension_name, diff_important_pos_unit_in_m_ / mm, int(numpy.ceil(
        diff_important_pos_unit_in_m_ / size_of_unimportant_cell))) for diff_important_pos_unit_in_m_ in
                     diff_important_pos_unit_in_m])
    copyrightinfo_end = "\n! =========auto-generated code end=========\n"

    cmds += copyrightinfo_end
    logger.info(cmds)
    return cmds


if __name__ == '__main__':
    plt.ion()
    sldprts = []

    parent_sldprt = r"E:\CSTprojects\GeneratorAccelerator\geom\Genac2\ModerateBuncher\temp\CoaxialRFsource.ModerateBuncher3.SideCoupling.ForMAGIC.SLDPRT"
    # parent_sldprt = r"E:\CSTprojects\dual_beam_dual_freq_HPM\egun\geom\yyh\egun.yyh.SLDPRT"
    # parent_sldprt = r"E:\CSTprojects\dual_beam_dual_freq_HPM\egun\geom\egun.SLDPRT"
    # parent_sldprt = r"D:\MagicFiles\HPM\电子枪部分\常规\电极形状20230720起.SLDPRT"
    sldprts_dir, parent_sldprt_filename = os.path.split(parent_sldprt)
    for file in os.listdir(sldprts_dir):
        if file.endswith(parent_sldprt_filename): continue  # 母文件
        if file.endswith('.SLDPRT') and not file.startswith("~$"):
            sldprts.append(os.path.join(sldprts_dir, file))
    logger.info(sldprts)

    # sldprts = [
    #     r"D:\MagicFiles\HPM\12.5GHz\swCATHODE.sldprt",
    #     r"D:\MagicFiles\HPM\12.5GHz\swANODE.sldprt",
    #     r"D:\MagicFiles\HPM\12.5GHz\swEMITTER.sldprt",
    #     r"D:\MagicFiles\HPM\12.5GHz\swPORT_LEFT.sldprt",
    #     r"D:\MagicFiles\HPM\12.5GHz\swPORT_RIGHT.sldprt"
    # ]

    # stl_files = [r"D:\MagicFiles\HPM\12.5GHz\swCATHODE.STL",
    #              r"D:\MagicFiles\HPM\12.5GHz\swANODE.STL",
    #              r"D:\MagicFiles\HPM\12.5GHz\swEMITTER.STL",
    #              r"D:\MagicFiles\HPM\12.5GHz\swPORT_LEFT.STL",
    #              r"D:\MagicFiles\HPM\12.5GHz\swPORT_RIGHT.STL"
    #              ]
    # stl_files = [r"D:\MagicFiles\HPM\电子枪部分\阴极附近陶瓷\swANODE.STL",
    #              r"D:\MagicFiles\HPM\电子枪部分\阴极附近陶瓷\swCATHODE.STL",
    #              r"D:\MagicFiles\HPM\电子枪部分\阴极附近陶瓷\swEMITTER.STL",
    #              r"D:\MagicFiles\HPM\电子枪部分\阴极附近陶瓷\swPORT_LEFT.STL",
    #              r"D:\MagicFiles\HPM\电子枪部分\阴极附近陶瓷\swPORT_RIGHT.STL"
    #              ]

    cmd = sldprt_to_MAGIC_commands(sldprts, axis_of_sketch=(Axis.z, Axis.y), )
    temp_file_to_save_MAGIC_comands = 'temp.txt'
    with open(temp_file_to_save_MAGIC_comands, 'w') as f:
        f.write(cmd)
        logger.info(cmd)
    os.system('code %s' % temp_file_to_save_MAGIC_comands)
    plt.legend()
    # plt.gca().set_aspect('equal')
    plt.show()

    MAGIC_object_name_prefixs = legend_txt = [t.get_text() for t in plt.legend().get_texts()]
    logger.info("Those commands might be also useful:\n"
                + "\n" * 3
                + "\n".join(
        ["VOID %s.0;" % MAGIC_object_name_prefix for MAGIC_object_name_prefix in MAGIC_object_name_prefixs])
                + "\n" * 3
                + "\n".join(["MARK %s.0 x1 size 0.5mm;" % MAGIC_object_name_prefix for MAGIC_object_name_prefix in
                             MAGIC_object_name_prefixs])
                + "\n" * 3
                + "\n".join(["MARK %s.0 x2 size 0.5mm;" % MAGIC_object_name_prefix for MAGIC_object_name_prefix in
                             MAGIC_object_name_prefixs])
                )
    plt.savefig("geom_view.png")
